import torch
import torch.nn as nn
import numpy as np
torch.manual_seed(1)

feats = torch.rand((2,3,5))
a=torch.full(feats.size(),1)

print(a)
print(feats,feats.size())
print("feats[:,0,:]",feats[:,0,:])

test = feats.transpose(0,1)


print("test", test,test.size())
print("test[0]",test[0])

print("===========================",torch.equal(feats[:,0,:],(test[0])))
transision = torch.Tensor(np.random.random((5,5)))

print("transision:",transision,transision[0])
print("***********"*10)
START_TAG = "<START>"
STOP_TAG = "<STOP>"

tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}
transision[:, tag_to_ix[STOP_TAG]] = -1000
transision[tag_to_ix[START_TAG]:,] = -1000
batch_size = feats.size()[0]
init_vvars = torch.full((batch_size, len(tag_to_ix)), -10000.)
init_vvars[:,tag_to_ix[START_TAG]] = 0
a = transision[0]
# print(type(a),a.size(),a)
# forward_var at step i holds the viterbi variables for step i-1
forward_var = init_vvars
print(forward_var)
result = forward_var+a
print(result)
best_ids = torch.argmax(forward_var,dim=1).view(-1,1)
print(best_ids,type(best_ids),best_ids.size())
ad = result.gather(1,best_ids)
print(ad)
backpointers = []





for feat in feats.transpose(0,1):
    bptrs_t = []  # holds the backpointers for this step
    viterbivars_t = []
    for nex_tag in range(len(tag_to_ix)):
        next_tag_var = forward_var+transision[nex_tag]
        best_ids = torch.argmax(forward_var,dim=1).unsqueeze(1)
        bptrs_t.append(best_ids)
        best_node_score = torch.gather(next_tag_var,dim=1,index=best_ids)
        viterbivars_t.append(best_node_score)
    forward_var = torch.cat(viterbivars_t,1)+feat
    backpointers.append(torch.cat(bptrs_t, 1))

terminal_var = forward_var+transision[tag_to_ix[STOP_TAG]]
best_tag_id = torch.argmax(forward_var,dim=1).unsqueeze(1)
path_score = torch.gather(terminal_var,dim=1,index=best_tag_id)
print("path_score:",path_score)
print(backpointers)

best_path = [best_tag_id]

for path in reversed(backpointers):
    best_tag_id = torch.gather(path, dim=1, index=best_tag_id)
    best_path.append(best_tag_id)
first_tag = best_path.pop()
print(first_tag)
assert first_tag.equal(torch.full_like(first_tag,3))
best_path.reverse()
print(torch.cat(best_path,1))
